Add support for unified hosts by jh-db · Pull Request #1307 · databricks/databricks-sdk-go

jh-db · 2025-10-15T09:35:03Z

What changes are proposed in this pull request?

This PR adds support for unified hosts, i.e. hosts that support both workspace and account methods. To do this, it:

separates client type from host type determination, deprecating isAccountClient and replacing it with HostType and ConfigType.
adds an experimental flag to indicate if a host is unified: experimental_is_unified_host
adds a WorkspaceId attribute to config, which is necessary for workspace clients that talk to unified hosts
adds UnifiedOAuthArgument, which is used in endpoint_supplier.go to discover OAuth endpoints on unified hosts and cache tokens for those hosts.
Adds a request visitor which adds an X-Databricks-Org-Id header to requests made by workspace clients

Note: workspace_client.go and account_client.go are generated and will be updated in a separate PR upstream.

How is this tested?

unit & intg tests

mgyucht · 2025-10-24T12:30:31Z

+// It relies on the workspace ID being set if and only if a workspace client is being used.
+func workspaceOrgIdVisitor(cfg *Config) func(r *http.Request) error {
+	return func(r *http.Request) error {
+		if cfg.WorkspaceId != "" {


suggestion: let's check what happens if a user specifies the wrong workspace ID for a given workspace host.

Just checked, the header is ignored by workspace hosts.

mgyucht · 2025-10-24T12:40:14Z

 	return a.ln.Close()
 }

 // validateArg ensures that the OAuthArgument is either a WorkspaceOAuthArgument


Update this comment

mgyucht · 2025-10-24T12:43:49Z

 }

 var ErrNotAccountClient = errors.New("invalid Databricks Account configuration")
+var ErrWorkspaceIdInAccountClient = errors.New("WorkspaceId must not be set when using AccountClient")


These changes need to happen in codegen.

will split these out, keeping them in here just for tests & review for now

mgyucht · 2025-10-24T12:48:04Z

 	}
-	if cfg.IsAccountClient() {
+	hostType := cfg.GetHostType()
+	if hostType == config.AccountHost || hostType == config.UnifiedHost && cfg.WorkspaceId == "" {


suggestion: separate check with a better error message than the (already poor) status quo

will do in codegen

mgyucht

Few comments/suggestions, let me know what you think

mgyucht · 2025-10-30T12:52:11Z

+		},
+	}
+
+	// Workspace IDs must be provided explicitly in request headers when using unified hosts.


suggestion: a bit more clarity on this doccomment, for example:

and when calling workspace-level APIs. Furthermore, workspace ID must not be provided when calling account-level APIs.

mgyucht · 2025-10-30T12:53:44Z

-		if err != nil {
-			return nil, fmt.Errorf("resolve host: %w", err)
-		}
+	err := cfg.azureEnsureWorkspaceUrl(ctx, c)


Stupid question: we can use Azure MSI with account-level APIs as well, right?

Yes, but we do various checks that basically ensure that azureEnsureWorkspaceUrl only takes effect if the client is configured for workspace APIs. But just to remove any doubt I've replaced the calls here and in auth_gcp_id with ClientType(). If there's a desire to remove these in the future we can build on the research done here, but we don't need to block this PR on it.

mgyucht · 2025-10-30T12:55:15Z

+// client in the future.
+//
+// Deprecated: Use HostType() instead.
+func (c *Config) ConfigType() ConfigType {


Where is this used?

It's needed by Terraform. We could use it in azure_msi and gcp_id auth.

mgyucht · 2025-10-30T12:55:20Z

+// ConfigType returns the type of config that the client is configured for.
+// Returns InvalidConfig if the config is invalid.
+// Use of this function should be avoided where possible, because we plan
+// to remove WorkspaceClient and AccountClient in favo of a single unified


Suggested change

// to remove WorkspaceClient and AccountClient in favo of a single unified

// to remove WorkspaceClient and AccountClient in favor of a single unified

mgyucht · 2025-10-30T12:57:18Z

+// to remove WorkspaceClient and AccountClient in favo of a single unified
+// client in the future.
+//
+// Deprecated: Use HostType() instead.


I don't get the sense that this is deprecated yet, though I do agree that it would be deprecated once we move to the unified API endpoint. Until then, there isn't really an alternative to this, so it doesn't make sense to me to recommend using HostType() directly.

Yeah, I agree. I was trying to discourage usage, but I think marking it as deprecated is too much when there's no viable alternative yet in some cases.

mgyucht · 2025-10-30T12:58:36Z

+	}
 	host := c.CanonicalHostName()
-	if c.IsAccountClient() {
+	switch hostType := c.HostType(); hostType {


I think you can pattern match on the result of the assignment without repeating it in Golang

Suggested change

switch hostType := c.HostType(); hostType {

switch hostType := c.HostType() {

I'm using it in the error message, but I can just call the function twice if that's preferred?

mgyucht · 2025-10-30T12:58:59Z

 	}
 	host := c.CanonicalHostName()
-	if c.IsAccountClient() {
+	switch hostType := c.HostType(); hostType {


Ditto

Suggested change

switch hostType := c.HostType(); hostType {

switch hostType := c.HostType() {

mgyucht · 2025-10-30T13:03:50Z

+// UnifiedOAuthArgument is an interface that provides the necessary information
+// to authenticate using OAuth to a host that supports both account and workspace APIs.
+type UnifiedOAuthArgument interface {
+	OAuthArgument


Stupid question: this interface is identical to AccountOAuthArgument, can it just implement that interface instead?

mgyucht · 2025-10-30T14:25:17Z

+// UnifiedOAuthArgument is an interface that provides the necessary information
+// to authenticate using OAuth to a host that supports both account and workspace APIs.
+type UnifiedOAuthArgument interface {
+	OAuthArgument


For now, I don't see the use of having an extra interface. We can just document that the AccountOAuthArgument interface can be used whether authenticating to an accounts host or a unified host. If there were a net new behavior that only applied to unified hosts, we could add something like this. So let's remove this (and maybe touch up the documentation on the accounts interface/implementation).

mgyucht · 2025-10-30T14:29:36Z

 	case AccountOAuthArgument:
 		endpoints, err = a.endpointSupplier.GetAccountOAuthEndpoints(
 			a.ctx, argg.GetAccountHost(), argg.GetAccountId())
+	case UnifiedOAuthArgument:


OK, I see what I think is going a bit wrong here. I feel like the OAuthEndpointSupplier should actually be provided as part of the OAuthArgument. That way the PersistantAuth doesn't need to know which EndpointSupplier method to call for a given OAuth argument, and it is fixed per OAuth argument.

Agreed, though we can treat this as a follow-up.

renaudhartert-db · 2025-10-30T14:38:49Z

+	if c.AccountID != "" && c.isTesting {
+		return AccountHost
+	}


Suggested change

if c.AccountID != "" && c.isTesting {

return AccountHost

}

// TODO: Refactor tests so that this is not needed.

if c.AccountID != "" && c.isTesting {

return AccountHost

}

github-actions · 2025-10-30T22:04:46Z

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-go

Inputs:

PR number: 1307
Commit SHA: f5d779a046c445c01e3dadd0e00ad6ef8eb7090c

Checks will be approved automatically on success.

## What changes are proposed in this pull request? This PR adds support for unified host: - Separates client type from host type determination, deprecating is_account_client and replacing it with host_type and client_type properties using new HostType and ClientType enums - Adds an experimental flag to indicate if a host is unified: experimental_is_unified_host - Adds a workspace_id attribute to Config, which is necessary for workspace clients that talk to unified hosts - Adds get_unified_endpoints() function, which is used in the OIDC endpoint resolution logic to discover OAuth endpoints on unified hosts - Adds header injection logic in ApiClient which adds an X-Databricks-Org-Id header to requests made by workspace clients on unified hosts - Adds comprehensive test coverage including unit tests for host/config type detection, OIDC endpoint resolution, header injection, and integration tests Similar to what is done for databricks/databricks-sdk-go#1307 ## How is this tested? - Unit tests - Manually integration tested <img width="1275" height="331" alt="image" src="https://github.com/user-attachments/assets/f456e104-6210-43a4-a99a-1c2ff0b134a2" /> - All existing integration tests pass Note: Integration test would be added in another PR once the test infra supports spog

## Changes  - Add support for unified host with experimental flag. - Prompt for workspace id as it's not part of the host - Write the IDs and flag in the created profile. - Depends on databricks/databricks-sdk-go#1307 ## Why  This support is required for enabling unified hosts which Databricks free edition uses. ## Tests  - Unit tests - Manual E2E tests -- CUJ for CLI U2M `databricks auth login --host <spog-host> --experimental-is-unified-host` prompts for account and workspace id followed by opening web browser. The IDs and flag is stores in the config. The subsequent cli operations eg: databricks clusters list --profile "above-profile" works.

## What changes are proposed in this pull request? This PR adds support for unified host: - Separates client type from host type determination, deprecating `isAccountClient` and replacing it with `getHostType()` and `getClientType()` methods using new `HostType` and `ClientType` enums - Adds an experimental flag to indicate if a host is unified: `experimentalIsUnifiedHost` - Adds a `workspaceId` attribute to DatabricksConfig, which is necessary for workspace clients that talk to unified hosts - Adds `getUnifiedOidcEndpoints()` function, which is used in the OIDC endpoint resolution logic to discover OAuth endpoints on unified hosts - Adds header injection logic in DatabricksConfig.authenticate() which adds an `X-Databricks-Org-Id` header to requests made by workspace clients on unified hosts - Adds comprehensive test coverage including unit tests for host/client type detection, OIDC endpoint resolution, header injection, and integration tests Similar to what is done for: - [databricks/databricks-sdk-py#1135](databricks/databricks-sdk-py#1135) - [databricks/databricks-sdk-go#1307](databricks/databricks-sdk-go#1307) ## How is this tested? - Unit tests - Manually E2E tested with a spog profile

## Changes  - Add support for unified host with experimental flag. - Prompt for workspace id as it's not part of the host - Write the IDs and flag in the created profile. - Depends on databricks/databricks-sdk-go#1307 ## Why  This support is required for enabling unified hosts which Databricks free edition uses. ## Tests  - Unit tests - Manual E2E tests -- CUJ for CLI U2M `databricks auth login --host <spog-host> --experimental-is-unified-host` prompts for account and workspace id followed by opening web browser. The IDs and flag is stores in the config. The subsequent cli operations eg: databricks clusters list --profile "above-profile" works.

jh-db temporarily deployed to test-trigger-is October 15, 2025 09:35 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 15, 2025 09:36 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 15, 2025 09:37 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 16, 2025 07:50 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 16, 2025 07:51 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 16, 2025 07:53 — with GitHub Actions Inactive

jh-db added 6 commits October 16, 2025 11:11

support unified hosts

f878295

remove comment

e2a76fe

remove clientType mutation from config

63bb1c0

fix

f91efff

WIP?

2e61308

Bring back removed functions and errors

21f56ce

jh-db force-pushed the jh-db/unified-hosts branch from cf52930 to 21f56ce Compare October 23, 2025 08:46

jh-db temporarily deployed to test-trigger-is October 23, 2025 08:46 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 23, 2025 08:47 — with GitHub Actions Inactive

jh-db mentioned this pull request Oct 24, 2025

[WIP] Add experimental support for unified hosts databricks/cli#3817

Closed

mgyucht reviewed Oct 24, 2025

View reviewed changes

try to fix tests

78ae20e

jh-db temporarily deployed to test-trigger-is October 24, 2025 14:10 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 24, 2025 14:11 — with GitHub Actions Inactive

Address comments

35cef46

jh-db temporarily deployed to test-trigger-is October 24, 2025 15:18 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 24, 2025 15:19 — with GitHub Actions Inactive

more tests

f981aa1

jh-db temporarily deployed to test-trigger-is October 24, 2025 15:27 — with GitHub Actions Inactive

update comment

c6b84d0

jh-db requested a review from mgyucht October 27, 2025 13:00

without generated file changes

261dc29

jh-db temporarily deployed to test-trigger-is October 30, 2025 10:45 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 30, 2025 10:46 — with GitHub Actions Inactive

jh-db changed the title ~~[WIP] Add support for unified hosts~~ Add support for unified hosts Oct 30, 2025

mgyucht reviewed Oct 30, 2025

View reviewed changes

address comments

97139a8

jh-db temporarily deployed to test-trigger-is October 30, 2025 14:13 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 30, 2025 14:14 — with GitHub Actions Inactive

mgyucht reviewed Oct 30, 2025

View reviewed changes

renaudhartert-db reviewed Oct 30, 2025

View reviewed changes

renaudhartert-db approved these changes Oct 30, 2025

View reviewed changes

changelog

b6b6c8c

jh-db temporarily deployed to test-trigger-is October 30, 2025 21:59 — with GitHub Actions Inactive

jh-db temporarily deployed to test-trigger-is October 30, 2025 22:00 — with GitHub Actions Inactive

Merge branch 'main' into jh-db/unified-hosts

f5d779a

jh-db temporarily deployed to test-trigger-is October 30, 2025 22:04 — with GitHub Actions Inactive

jh-db added this pull request to the merge queue Oct 31, 2025

Merged via the queue into main with commit 5b291af Oct 31, 2025
15 checks passed

jh-db deleted the jh-db/unified-hosts branch October 31, 2025 07:44

tanmay-db mentioned this pull request Dec 5, 2025

[Feature] Add support for Unified Host with experimental flag databricks/databricks-sdk-py#1135

Merged

tanmay-db mentioned this pull request Jan 19, 2026

[Feature] Support for Unified Host databricks/databricks-sdk-java#612

Merged

denik mentioned this pull request Jan 19, 2026

[Feature] Add support for Unified Host with experimental flag databricks/cli#4260

Merged

tanmay-db mentioned this pull request Feb 2, 2026

[Internal] Refactor WorkspaceId to WorkspaceID in config #1457

Closed

	// to remove WorkspaceClient and AccountClient in favo of a single unified
	// to remove WorkspaceClient and AccountClient in favor of a single unified

	switch hostType := c.HostType(); hostType {
	switch hostType := c.HostType() {

Uh oh!

Conversation

jh-db commented Oct 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes are proposed in this pull request?

How is this tested?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mgyucht left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jh-db Oct 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 30, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jh-db commented Oct 15, 2025 •

edited

Loading

jh-db Oct 30, 2025 •

edited

Loading